package com.ls.rocketmq.agent;

import net.bytebuddy.agent.builder.AgentBuilder;

import java.lang.instrument.Instrumentation;
import java.util.HashMap;
import java.util.Map;

import static net.bytebuddy.matcher.ElementMatchers.named;

public class RocketMqTopicIpReplacerAgent {
    public static final Map<String, String> IP_MAPPING = new HashMap<>();

    public static void premain(String arguments, Instrumentation instrumentation) {
        System.out.println("[RocketMQ Topic Ip Replacer Agent] 进入修改class ");
        System.out.println("[RocketMQ Topic Ip Replacer Agent] 参数: " + arguments);
        parseAgentArgs(arguments);

        new AgentBuilder.Default()
            .type(named("org.apache.rocketmq.client.impl.MQClientAPIImpl"))
            .transform(new MQClientAPIImplTransformer())
            .installOn(instrumentation);
    }

    private static void parseAgentArgs(String agentArgs) {
        if (agentArgs != null) {
            String[] mappings = agentArgs.split(",");
            for (String mapping : mappings) {
                String[] ips = mapping.split(":");
                if (ips.length == 2) {
                    IP_MAPPING.put(ips[0].trim(), ips[1].trim());
                    System.out.println("[RocketMQ Topic Ip Replacer Agent] Added IP mapping: " + ips[0] + " -> " + ips[1]);
                }
            }
        }
    }
}
